import math
import os

import torch
import torch.nn as nn
import torch.utils.model_zoo as model_zoo

BatchNorm2d = nn.BatchNorm2d


# 普通Conv3x3卷积块
def conv_bn(inp, oup, stride):
    return nn.Sequential(
        nn.Conv2d(
            in_channels=inp,
            out_channels=oup,
            kernel_size=3,
            stride=stride,
            padding=1,
            bias=False,
        ),
        BatchNorm2d(num_features=oup),
        nn.ReLU6(inplace=True),
    )


# 普通Conv1x1卷积块
def conv_1x1_bn(inp, oup):
    return nn.Sequential(
        nn.Conv2d(
            in_channels=inp,
            out_channels=oup,
            kernel_size=1,
            stride=1,
            padding=0,
            bias=False,
        ),
        BatchNorm2d(num_features=oup),
        nn.ReLU6(inplace=True),
    )


class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = round(inp * expand_ratio)
        # shortcut连接
        self.use_res_connect = self.stride == 1 and inp == oup

        if expand_ratio == 1:
            self.conv = nn.Sequential(
                # --------------------------------------------#
                #   进行3x3的逐层卷积，进行跨特征点的特征提取
                # --------------------------------------------#
                nn.Conv2d(
                    in_channels=hidden_dim,
                    out_channels=hidden_dim,
                    kernel_size=3,
                    stride=stride,
                    padding=1,
                    groups=hidden_dim,
                    bias=False,
                ),
                BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                # -----------------------------------#
                #   利用1x1卷积进行通道数的调整
                # -----------------------------------#
                nn.Conv2d(
                    in_channels=hidden_dim,
                    out_channels=oup,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                    bias=False,
                ),
                BatchNorm2d(oup),
            )
        else:
            self.conv = nn.Sequential(
                # -----------------------------------#
                #   利用1x1卷积进行通道数的上升
                # -----------------------------------#
                nn.Conv2d(
                    in_channels=inp,
                    out_channels=hidden_dim,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                    bias=False,
                ),
                BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                # --------------------------------------------#
                #   进行3x3的逐层卷积，进行跨特征点的特征提取（使用组卷积）
                # --------------------------------------------#
                nn.Conv2d(
                    in_channels=hidden_dim,
                    out_channels=hidden_dim,
                    kernel_size=3,
                    stride=stride,
                    padding=1,
                    groups=hidden_dim,
                    bias=False,
                ),
                BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                # -----------------------------------#
                #   利用1x1卷积进行通道数的下降
                # -----------------------------------#
                nn.Conv2d(
                    in_channels=hidden_dim,
                    out_channels=oup,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                    bias=False,
                ),
                BatchNorm2d(oup),
            )

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


class MobileNetV2(nn.Module):
    def __init__(self, n_class=1000, input_size=224, width_mult=1.0):
        super(MobileNetV2, self).__init__()
        block = InvertedResidual
        input_channel = 32
        last_channel = 1280
        interverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],  # 256, 256, 32 -> 256, 256, 16
            [6, 24, 2, 2],  # 256, 256, 16 -> 128, 128, 24   2
            [6, 32, 3, 2],  # 128, 128, 24 -> 64, 64, 32     4
            [6, 64, 4, 2],  # 64, 64, 32 -> 32, 32, 64       7
            [6, 96, 3, 1],  # 32, 32, 64 -> 32, 32, 96
            [6, 160, 3, 2],  # 32, 32, 96 -> 16, 16, 160     14
            [6, 320, 1, 1],  # 16, 16, 160 -> 16, 16, 320
        ]

        assert input_size % 32 == 0
        input_channel = int(input_channel * width_mult)
        self.last_channel = (
            int(last_channel * width_mult) if width_mult > 1.0 else last_channel
        )
        # (512, 512, 3) -> (256, 256, 32)
        self.features = [conv_bn(inp=3, oup=input_channel, stride=2)]

        for t, c, n, s in interverted_residual_setting:
            # t - expansion  c - output channel  n - number of bottleneck  s - stride
            output_channel = int(c * width_mult)
            for i in range(n):
                if i == 0:
                    self.features.append(
                        block(input_channel, output_channel, stride=s, expand_ratio=t)
                    )  # 仅在第一层进行下采样
                else:
                    self.features.append(
                        block(input_channel, output_channel, stride=1, expand_ratio=t)
                    )  # bottleneck模块中后续卷积层不改变输入的(h, w)
                input_channel = output_channel  # 更新input_channel

        self.features.append(
            conv_1x1_bn(inp=input_channel, oup=self.last_channel)
        )  # 320 -> 1280
        self.features = nn.Sequential(*self.features)

        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(in_features=self.last_channel, out_features=n_class),
        )
        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = x.mean(3).mean(2)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                n = m.weight.size(1)
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()


def load_url(url, model_dir="./model_data", map_location=None):
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
    filename = url.split("/")[-1]
    cached_file = os.path.join(model_dir, filename)
    if os.path.exists(cached_file):
        return torch.load(cached_file, map_location=map_location)
    else:
        return model_zoo.load_url(url, model_dir=model_dir)


def mobilenetv2(pretrained=False, **kwargs):
    model = MobileNetV2(n_class=1000, **kwargs)
    if pretrained:
        model.load_state_dict(torch.load("./model_data/mobilenet_v2.pth"))
        '''        
        model.load_state_dict(
            load_url(
                "https://github.com/bubbliiiing/deeplabv3-plus-pytorch/releases/download/v1.0/mobilenet_v2.pth.tar"
            ),
            strict=False,
        )
        '''

    return model


if __name__ == "__main__":
    model = mobilenetv2()
    for i, layer in enumerate(model.features):
        print(i, layer)
